#!/usr/bin/env python

import socket

port = 8000

def get_host(request_data):
    string = request_data.decode()
    host = ''
    i = 11
    while True:
        byte = string[i]
        if byte == ' ' or byte == '/':
            break
        host += byte
        i += 1
    print(f'Host: {host}')
    return host

def forward_request(host, request_data):
    remote_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    remote_socket.connect((host, 80))
    remote_socket.sendall(request_data)
    print('Forwarded request to a remote server')
    return remote_socket

def get_headers(remote_socket):
    response_data = b''
    flag = False
    while True:
        byte = remote_socket.recv(1)
        response_data += byte
        if response_data[-4:len(response_data)] == b'\r\n\r\n':
            break
    return response_data

def get_content_length(headers):
    i = headers.decode().find("Content-Length")
    print(headers.decode())
    flag = False
    string_num = ''
    while True:
        char = headers.decode()[i]
        if char == '\n':
            break
        if char == ' ':
            flag = True
        if not flag:
            i +=1
            continue
        string_num +=char
        i += 1
    return int(string_num)

def get_chunk_size(socket):
    chunk_size = b''
    while True:
        byte = socket.recv(1)
        if byte == b'\n':
            break
        chunk_size += byte
    #print("Size: ", chunk_size, " | ", int(chunk_size.decode(), 16) ," <<")
    return int(chunk_size.decode(), 16)

def get_chunks(socket):
    response_data = b''
    while True:
        chunk_size = get_chunk_size(socket)
        if chunk_size == 0:
            break
        response_data += socket.recv(chunk_size+2)
    print(response_data.decode())
    return response_data

def get_response_data(socket, headers):
    if headers.decode().find('chunked') != -1:
        response_data = get_chunks(socket)
    else: 
        content_length = get_content_length(headers)
        response_data = socket.recv(content_length)
    return response_data


def run_proxy_server(port):
    print("Start proxy server")
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(('localhost', port))
    server_socket.listen()

    while True:
        print(f'Proxy listen {port}...')
        client_socket, client_address = server_socket.accept()
        print(f'Connect to client: {client_address}')
        
        request_data = client_socket.recv(4096)
        host = get_host(request_data)
        try:
            remote_socket = forward_request(host, request_data)
        except:
            continue
        headers = get_headers(remote_socket)
        print(headers.decode())

        response_data = get_response_data(remote_socket, headers) 
        
        print(f'Got response:\n{response_data}')
        client_socket.sendall(response_data)
        client_socket.close()

if __name__ == '__main__':
    run_proxy_server(port)

